import os
import json
import pymysql
from dotenv import load_dotenv
from mcp.server.fastmcp import FastMCP
from typing import Optional

# 加载环境变量
load_dotenv(override=True)

# MySQL连接配置
MYSQL_CONFIG = {
    'host': os.getenv('MYSQL_HOST', 'localhost'),
    'port': int(os.getenv('MYSQL_PORT', 3306)),
    'user': os.getenv('MYSQL_USER', 'root'),
    'password': os.getenv('MYSQL_PASSWORD', ''),
    'database': os.getenv('MYSQL_DATABASE', 'user_insight'),
    'charset': 'utf8mb4',
    'cursorclass': pymysql.cursors.DictCursor  # 使结果自动转为字典格式
}

# 创建MCP
mcp = FastMCP("AgentMySQLServer")


@mcp.tool()
async def sql_inquire(sql_query: str) -> str:
    """
    当用户需要进行数据库查询工作时，请调用该函数。
    该函数用于在MySQL数据库上运行SQL查询语句。

    参数:
        sql_query (str): SQL查询语句

    返回:
        str: JSON格式的查询结果或错误信息
    """
    connection = None
    try:
        # 创建MySQL连接
        connection = pymysql.connect(**MYSQL_CONFIG)

        with connection.cursor() as cursor:
            cursor.execute(sql_query)

            # 获取查询结果
            if sql_query.strip().upper().startswith('SELECT'):
                results = cursor.fetchall()
                return json.dumps(results, ensure_ascii=False)
            else:
                # 对于非SELECT语句，返回影响的行数
                connection.commit()
                return json.dumps({"affected_rows": cursor.rowcount}, ensure_ascii=False)

    except Exception as e:
        # 回滚事务并返回错误信息
        if connection:
            connection.rollback()
        return json.dumps({"error": str(e)}, ensure_ascii=False)
    finally:
        if connection:
            connection.close()



if __name__ == "__main__":
    # 测试连接
    try:
        test_conn = pymysql.connect(**MYSQL_CONFIG)
        print("✅ MySQL连接成功")
        test_conn.close()
    except Exception as e:
        print(f"❌ MySQL连接失败: {e}")
        exit(1)

    mcp.run(transport="stdio")